#!/usr/bin/env python
"""
Suite of microbenchmarks.
"""
__author__ = "Dan Gunter <dkgunter@lbl.gov>"
__rcsid__ = "$Id$"

import math
import optparse
import os
import re
import sys
import time
#
from amqplib import client_0_8 as amqp

# Names of benchmark types
T_MULTI = 'multi'
T_SIZE = 'size'

class Benchmark:
    POW = 2 # powers of this define the size increments

    def __init__(self):
        pass

    def run_all(self, iterations, **kw):
        self.run_num, self.final = 0, False
        for self.run_num in range(iterations):
            if self.run_num == iterations - 1:
                self.final = True
            self.run(**kw)
        return 0


    def pow2_range(self, minsz, maxsz):
        """Build and return list of powers of 2
        between minsz and maxsz, inclusive.
        """
        # special case where range is just one item
        if minsz == maxsz:
            return [minsz]
        sizes = [ ]
        n = minsz
        while n < maxsz:
            sizes.append(n)
            n = n * self.POW
        sizes.append(maxsz)
        return sizes

    def run(self, mode=None, host='localhost', nmsg=1):
        """Run in a given (string-defined) mode, with
        the given RabbitMQ host and number of messages
        """
        pass

class SizeBenchmark(Benchmark):
    OUTPUT_HDR = "run,size,dur,count"
    EXCHANGE = 'msgsize_ex'
    QUEUE = 'msgsize_q'


    def __init__(self, range_=(1,128)):
        Benchmark.__init__(self)
        self.sizes = self.pow2_range(*range_)

    def run(self, mode=None, host='localhost', nmsg=1):
        channel = self.connect(host)
        if mode == 'p':
            if self.run_num == 0:
                print self.OUTPUT_HDR
            for sz in self.sizes:
                t0 = time.time()
                self.producer(channel, nmsg, size=sz)
                t1 = time.time()
                g = t1 - t0
                self.writerow(sz, g, nmsg)
        elif mode == 'c':
            self.consumer(channel, nmsg * len(self.sizes))
        channel.close()

    def writerow(self, sz, tm, count):
        print "%d,%d,%f,%d" % (self.run_num+1, sz, tm, count)

    def connect(self, host):
        conn = amqp.Connection(host=host)
        channel = conn.channel()
        channel.exchange_declare(exchange=self.EXCHANGE, type='direct',
                                 durable=True, auto_delete=False)
        channel.queue_declare(queue=self.QUEUE, durable=True, auto_delete=False)
        channel.queue_bind(queue=self.QUEUE, exchange=self.EXCHANGE, routing_key=self.QUEUE)
        return channel

    def producer(self, channel, n, size=1):
        buf = chr(0) * size
        for i in xrange(n):
            msg = amqp.Message(buf, delivery_mode=2)
            channel.basic_publish(msg, exchange=self.EXCHANGE, routing_key=self.QUEUE,
                                  mandatory=True)

    def consumer(self, channel, n):
        counter = Counter(channel, n)
        tag = channel.basic_consume(queue=self.QUEUE, callback=counter.got_msg)
        while counter.remain > 0:
            channel.wait()
        channel.basic_cancel(tag)

class MultiBenchmark(Benchmark):

    OUTPUT_HDR = "run,msgsize,dur,numq,count"
    EXCHANGE = 'multi_ex'
    QUEUE = 'multi_q'
    RESULT_QUEUE = 'multi_result_q'
    DONE_RUN = "DONE1"
    DONE_ALL = "DONE"

    def __init__(self, queues=1, rank=None, msg_size=0):
        Benchmark.__init__(self)
        self.queues = queues
        self.rank = rank
        self.msg_size = msg_size

    def connect(self, host):
        conn = amqp.Connection(host=host)
        self.channel = conn.channel()
        self.channel.exchange_declare(exchange=self.EXCHANGE, type='direct',
                                      durable=True, auto_delete=False)

        # result queue
        self.channel.queue_declare(queue=self.RESULT_QUEUE, durable=True, auto_delete=False)
        self.channel.queue_bind(queue=self.RESULT_QUEUE, exchange=self.EXCHANGE,
                                routing_key=self.RESULT_QUEUE)

    def run(self, mode=None, host='localhost', nmsg=1):
        self.connect(host)
        if mode == 'p':
            if self.run_num == 0:
                print self.OUTPUT_HDR
            # setup
            self.create_queues(0, self.queues)
            # send messages
            t0 = time.time()
            self.producer(nmsg, self.queues)
            # read timestamps from clients
            t1 = 0
            for i in xrange(self.queues):
                while 1:
                    msg = self.channel.basic_get(queue=self.RESULT_QUEUE, no_ack=True)
                    if msg: break
                    time.sleep(1)
                ts = float(msg.body)
                t1 = max(t1, ts)
            # calculate total time and write it 
            deltat = t1 - t0
            self.writerow(self.msg_size, deltat, self.queues, nmsg)
            # close channel
            if self.final:
                self.channel.close()
        elif mode == 'c':
            self.consumer_done = False
            self.create_queues(0, self.queues)
            self.consumer()
        return 0

    def get_qname(self, num):
        return self.QUEUE + '_' + str(num)

    def create_queues(self, q1, q2):
        for q in range(q1, q2):
            qname = self.get_qname(q)
            #print "create queue %s"  % qname
            self.channel.queue_declare(queue=qname, durable=True, auto_delete=False)
            self.channel.queue_bind(queue=qname, exchange=self.EXCHANGE,
                               routing_key=qname)

    def producer(self, nmsg, qnum):
        """Run producer
        """
        buf = chr(0) * self.msg_size
        qnames = map(self.get_qname, range(qnum))
        # send messages
        for i in xrange(nmsg):
            msg = amqp.Message(buf, delivery_mode=2)
            # choose queues round-robin
            qname = qnames[i % qnum]
            # send msg to chosen queue
            self.channel.basic_publish(msg, exchange=self.EXCHANGE,
                                  routing_key=qname, mandatory=False,
                                  immediate=True)
        # send 'final' message for this run
        msg = amqp.Message(self.DONE_RUN, delivery_mode=2)
        for q in xrange(qnum):
            qname = self.get_qname(q)
            # send msg to chosen queue
            self.channel.basic_publish(msg, exchange=self.EXCHANGE,
                                       routing_key=qname, mandatory=True)
        # if last run, tell consumers to shutdown
        if self.final:
            msg = amqp.Message(self.DONE_ALL, delivery_mode=2)
            for q in xrange(qnum):
                qname = self.get_qname(q)
                # send msg to chosen queue
                self.channel.basic_publish(msg, exchange=self.EXCHANGE,
                                      routing_key=qname, mandatory=True)

    def consumer_cb(self, msg):
        """Process message for consumer.
        """
        self.channel.basic_ack(msg.delivery_tag)
        if msg.body == self.DONE_RUN:
            self.consumer_send_timestamp(time.time())
        elif msg.body == self.DONE_ALL:
            self.channel.basic_cancel(self.consumer_tag)
            self.consumer_done = True
            time.sleep(2)
            self.channel.close()
            sys.exit(0) # hack! otherwise, blocks on read()

    def consumer_send_timestamp(self, t):
        """Send a timestamp back to the producer.
        """
        msg = amqp.Message("%f" % t, delivery_mode=2)
        self.channel.basic_publish(msg, exchange=self.EXCHANGE, routing_key=self.RESULT_QUEUE)
        
    def consumer(self):
        """Run consumer.
        """
        qname = self.get_qname(self.rank)
        self.consumer_tag = self.channel.basic_consume(queue=qname,
                                                       callback=self.consumer_cb)
        while not self.consumer_done:
            self.channel.wait()
                                   
    def writerow(self, sz, tm, numq, count):
        print "%d,%d,%f,%d,%d" % (self.run_num+1, sz, tm, numq, count)

class Counter:
    def __init__(self, channel, remain):
        self.remain = remain
        self.channel = channel

    def got_msg(self, m):
        self.channel.basic_ack(m.delivery_tag)
        self.remain -= 1

def parse_range(value):
    """Parse an integer range in the format 'a:b'
    Where both must be positive, b > a.
    Both a and b are rounded to the nearest power of 2.
    """
    m = re.match("(\d+):(\d+)$", value)
    if m is None:
        raise ValueError("Bad size, must be in format a:b")
    startsz, endsz = map(int, m.groups())
    if startsz < 0 or startsz > endsz:
        raise ValueError("Bad values for range a:b, "
                         "must be positive with a < b")
    # round to nearest power of 2
    round2 = lambda x:  int(pow(2,round(math.log(x,2))))
    return map(round2, (startsz, endsz))

def range_option_cb(option, opt, value, parser):
    try:
        value2 = parse_range(value)
    except ValueError, err:
        msg = str(err)
        raise optparse.OptionValueError(msg)
    setattr(parser.values, option.dest, value2)

def main(arguments=None):
    if arguments is None:
        cmdline = sys.argv[1:]
    else:
        cmdline = arguments
    desc = ' '.join(__doc__.split())
    opp = optparse.OptionParser(usage="%prog TYPE MODE HOST [options..]\n"
                                "  TYPE = <name>, where name is one of:\n"
                                "    size = Single producer, single consumer, "
                                "increasing message-size\n"
                                "    multi = Single producer, increasing number"
                                "of consumers/queues\n"
                                "  MODE = c for consumer, p for producer\n"
                                "  HOST = RabbitMQ server host" ,
                                description=desc)
    # General options
    grp = optparse.OptionGroup(opp, "General options")
    grp.add_option("-c", "--count", dest="count", type="int", default=100,
                   metavar="NUM", help="Number of messages (default=%default)")
    grp.add_option("-i", "--iter", dest="iter", type="int", default=10,
                   metavar="NUM",
                   help="Number of iterations to run (default=%default)")
    opp.add_option_group(grp)
    # Options for benchmark type = size
    grp = optparse.OptionGroup(opp, "TYPE = %s" % T_SIZE)
    grp.add_option("--sizes", action="callback", callback=range_option_cb,
                   dest='sizerange', default="512:2097152", metavar="RANGE",
                   help="Message size range; rounded to powers of 2 "
                   "(default=%default)")
    opp.add_option_group(grp)
    # Options for benchmark type = multi
    grp = optparse.OptionGroup(opp, "TYPE = %s" % T_MULTI)
    grp.add_option("--queues", type="int",
                   default=1, metavar="NUM", dest='queues',
                   help="Number of queues & consumers; rounded to powers of 2 "
                   "(default=%default)")
    grp.add_option("--rank", type="int", dest="rank", metavar="NUM",
                   default=None, help="For consumers, rank 0 .. N-1 of "
                   "this instance")
    opp.add_option_group(grp)
    grp.add_option("--size", type="int",
                   dest='msg_size', default=1024*1024, metavar="BYTES",
                   help="Message size (default=%default)")
    # Parser options
    options, args = opp.parse_args(cmdline)
    if len(args) != 3:
        opp.error("Wrong number of arguments")
    type_, mode, host = args[0].lower(), args[1], args[2]
    n = options.count
    # Branch for different benchmarks
    if type_ == T_SIZE:
        if isinstance(options.sizerange, str):
            # option not given, parse default value
            options.sizerange = parse_range(options.sizerange)
        bmark = SizeBenchmark(options.sizerange)
    elif type_ == T_MULTI:
        if mode == 'c' and options.rank is None:
            opp.error("Rank is required for consumers")
        bmark = MultiBenchmark(queues=options.queues, rank=options.rank,
                msg_size=options.msg_size)
    else:
        opp.error("Unknown benchmark TYPE (%s)" % type_)
    return bmark.run_all(options.iter, mode=mode, host=host, nmsg=n)

if __name__ == '__main__':
   sys.exit(main())
